package com.shen.langchain4j.config;

import com.shen.langchain4j.service.ChatAssistant;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * 大模型配置类
 */
@Configuration
public class LLMConfig {
    @Bean
    public ChatModel chatModel() {
        return OpenAiChatModel.builder()
                .apiKey(System.getenv("aliyunQwen-apiKey"))
                .modelName("qwen3-next-80b-a3b-instruct")
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .build();
    }

    @Bean
    public EmbeddingModel embeddingModel() {
        return OpenAiEmbeddingModel.builder()
                .apiKey(System.getenv("aliyunQwen-apiKey"))
                .modelName("text-embedding-v4")
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .build();
    }

    /**
     * 构建向量数据库客户端实例
     *
     * @return 向量客户端实例
     */
    @Bean
    public QdrantClient qdrantClient() {
        QdrantGrpcClient.Builder clientBuilder = QdrantGrpcClient
                .newBuilder("47.115.218.172", 6334, false);
        return new QdrantClient(clientBuilder.build());
    }

    /**
     * 基于Qdrant向量数据库的向量存储
     *
     * @return 嵌入存储
     */
    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return QdrantEmbeddingStore.builder()
                .host("47.115.218.172")
                .port(6334)
                .collectionName("test-qdrant")
                .build();
    }

    /**
     * 基于内存向量数据库的向量存储
     *
     * @return 嵌入存储
     */
    @Bean("memoryEmbeddingStroe")
    public InMemoryEmbeddingStore<TextSegment> inMemoryEmbeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    /**
     * 构建带聊天记忆，RAG功能的AI服务实例
     *
     * @param chatModel              大模型实例
     * @param inMemoryEmbeddingStore 内存向量数据库嵌入存储
     * @return AI服务实例
     */
    @Bean
    public ChatAssistant chatAssistant(ChatModel chatModel, @Qualifier("memoryEmbeddingStroe") EmbeddingStore<TextSegment> inMemoryEmbeddingStore) {
        return AiServices.builder(ChatAssistant.class)
                .chatModel(chatModel)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(100))
                .contentRetriever(EmbeddingStoreContentRetriever.from(inMemoryEmbeddingStore))
                .build();
    }
}
